# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

from coin.base.book import BookBuilderBase
from coin.base.book.aggregated_book_builder import BookUpdateSide
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct


class BookEntry(object):
  def __init__(self, id, price, qty, side):
    self.id = id
    self.price = float(price)
    self.qty = float(qty)
    self.side = side


class BookBuilder(BookBuilderBase):
  ProductType = BitmexFuturesProduct

  def __init__(self, native_symbol, book_builder_observer):
    super().__init__("Bitmex", native_symbol, book_builder_observer=book_builder_observer)
    self._id_map = {}

  def get_insert_diff(self, data: dict):
    price = data['price']
    qty = data['size']
    side = data['side']
    if data['side'] == 'Sell':
      return [BookUpdateSide.ASK, price, qty]
    else:
      return [BookUpdateSide.BID, price, qty]

  def get_update_diff(self, id: int, side: str, qty: int):
    if id not in self._id_map:
      return
    entry = self._id_map[id]
    price = entry.price
    old_qty = entry.qty
    if side == 'Sell':
      return [BookUpdateSide.ASK, price, qty - old_qty]
    else:
      return [BookUpdateSide.BID, price, qty - old_qty]

  def partial(self, data: dict):
    self.insert(data)

  def delete(self, id: int, side: str):
    if id not in self._id_map:
      return
    entry = self._id_map[id]
    price = entry.price
    del self._id_map[id]
    if side != entry.side:
      return
    if side == 'Sell':
      if price in self._ask_price_map:
        assert self._ask_price_map[price].id == id
        del self._ask_price_map[price]
    elif side == 'Buy':
      if price in self._bid_price_map:
        assert self._bid_price_map[price].id == id
        del self._bid_price_map[price]
    else:
      raise ValueError('Invalid side: %s' % side)

  def update(self, id: int, side: str, qty: int):
    if id not in self._id_map:
      return

    entry = self._id_map[id]
    price = entry.price
    entry.qty = qty

    if side != entry.side:
      if entry.side == 'Buy':
        if price in self._bid_price_map:
          del self._bid_price_map[price]
        self._ask_price_map[price] = entry
      elif entry.side == 'Sell':
        if price in self._ask_price_map:
          del self._ask_price_map[price]
        self._bid_price_map[price] = entry
      else:
        raise ValueError('Invalid side')
      entry.side = side

    else:
      # If the book is modified but PriceMap.__setitem__() is not called,
      # PriceMap.notify_updated() must be called explicitly.
      if entry.side == 'Buy':
        self._bid_price_map.notify_updated(price, entry)
      elif entry.side == 'Sell':
        self._ask_price_map.notify_updated(price, entry)
    self.remove_cross_book_level(entry)

  def insert(self, data: dict):
    assert data['size'] > 0, 'Size must be greater than 0.'
    id = data['id']
    price = data['price']
    if price is None:
      return

    assert price > 0
    qty = data['size']
    side = data['side']
    entry = BookEntry(id=id, price=price, qty=qty, side=side)
    if data['side'] == 'Sell':
      assert price not in self._ask_price_map or self._ask_price_map[price].id == id
      self._ask_price_map[price] = entry
    elif data['side'] == 'Buy':
      assert price not in self._bid_price_map or self._bid_price_map[price].id == id
      self._bid_price_map[price] = entry
    else:
      raise ValueError('Invalid side: %s' % data['side'])
    self._id_map[id] = entry
    self.remove_cross_book_level(entry)

  def remove_cross_book_level(self, entry):
    if entry.side == 'Sell':
      if self.bid0() is not None and entry.price <= self.bid0().price:
        self._delete_bid_up_to_and_including(entry.price)
    elif entry.side == 'Buy':
      if self.ask0() is not None and entry.price >= self.ask0().price:
        self._delete_ask_up_to_and_including(entry.price)
    else:
      raise ValueError('Invalid side: %s.' % entry.side)

  def _delete_bid_up_to_and_including(self, price):
    elist = self._bid_price_map.get_entries(price, float('Inf'), exclude_max=False)
    for entry in list(elist):
      assert entry.side == 'Buy'
      self.delete(entry.id, entry.side)

  def _delete_ask_up_to_and_including(self, price):
    elist = self._ask_price_map.get_entries(-float('Inf'), price, exclude_max=False)
    for entry in list(elist):
      assert entry.side == 'Sell'
      self.delete(entry.id, entry.side)
